import asyncio

from fastapi import FastAPI, Request

import bs4
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from langchain_core.vectorstores import InMemoryVectorStore
from langchain_chroma import Chroma
from sse_starlette.sse import EventSourceResponse
from langchain_core.messages import SystemMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

# 初始化 Vector Store
vector_store = Chroma(
    embedding_function=DashScopeEmbeddings(dashscope_api_key="sk-a3f7718fb81f43b2915f0a6483b6661b"),
    persist_directory="chroma_db"
)

# 工具定义
@tool(response_format="content_and_artifact")
def retrieve(query: str):
    """Retrieve information related to a query."""
    retrieved_docs = vector_store.similarity_search(query, k=1)
    serialized = "\n\n".join(
        (f"Source: {doc.metadata}\nContent: {doc.page_content}")
        for doc in retrieved_docs
    )
    return serialized, retrieved_docs

# LLM 实例
llm = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="llama-4-scout-17b-16e-instruct",
    streaming=True,
)

# Step 1: Tool call
def query_or_respond(state: MessagesState):
    llm_with_tools = llm.bind_tools([retrieve])
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

# Step 2: Execute tool
tools = ToolNode([retrieve])

# Step 3: Async generate with streaming
async def generate(state: MessagesState):
    recent_tool_messages = []
    for message in reversed(state["messages"]):
        if message.type == "tool":
            recent_tool_messages.append(message)
        else:
            break
    tool_messages = recent_tool_messages[::-1]

    docs_content = "\n\n".join(doc.content for doc in tool_messages)
    system_message_content = (
        "You are an assistant for question-answering tasks. "
        "Use the following pieces of retrieved context to answer "
        "the question. If you don't know the answer, say that you "
        "don't know. Use three sentences maximum and keep the "
        "answer concise."
        "\n\n"
        f"{docs_content}"
    )
    conversation_messages = [
        message
        for message in state["messages"]
        if message.type in ("human", "system")
        or (message.type == "ai" and not message.tool_calls)
    ]
    prompt = [SystemMessage(system_message_content)] + conversation_messages

    async for chunk in llm.astream(prompt):
        yield {"token": chunk.content}  # 发送单个 token
    # return {"messages": [chunk]}


# 构建图
graph_builder = StateGraph(MessagesState)
graph_builder.add_node("query_or_respond", query_or_respond)
graph_builder.add_node("tools", tools)
graph_builder.add_node("generate", generate)
graph_builder.set_entry_point("query_or_respond")
graph_builder.add_conditional_edges(
    "query_or_respond",
    tools_condition,
    {END: END, "tools": "tools"},
)
graph_builder.add_edge("tools", "generate")
graph_builder.add_edge("generate", END)
memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)


# FastAPI 应用
app = FastAPI()

@app.get("/stream")
async def stream(request: Request):
    query = request.query_params.get("query", "请讲一个关于猴子的笑话")
    config = {"configurable": {"thread_id": "abc123"}}

    async def event_generator():
        async for event in graph.astream(
            {"messages": [{"role": "user", "content": query}]},
            stream_mode="updates",
            config=config,
        ):
            for node_name, output in event.items():
                print(output.get("messages",[]))
                for message in output.get("messages", []):
                    print(message)
                    if hasattr(message, "type") and message.type == "ai":
                        if not getattr(message, "tool_calls", None):
                            yield {"data": message.content}
        yield {"data": "[DONE]"}

    return EventSourceResponse(event_generator())

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 启动命令：uvicorn test01:app --reload --host 0.0.0.0 --port 8000
